from PDE_Girkmann import Girkmann
from source_vector import source_vector
from set_essential_bc import set_essential_bc
from mass_vector_matrix import mass_vector_matrix
from compliance_tensor_matrix import compliance_tensor_matrix
from Post_processing import Post


import numpy as np
from scipy.sparse import spdiags, bmat
from scipy.sparse.linalg import spsolve
import argparse
import matplotlib.pyplot as plt


from fealpy.functionspace.HuZhangFiniteElementSpace2D import HuZhangFiniteElementSpace
from fealpy.functionspace.LagrangeFiniteElementSpace import LagrangeFiniteElementSpace




##  参数解析
parser = argparse.ArgumentParser(description=
        """
        三角形网格上用胡张元求解Girkmannproblem
        """)


parser.add_argument('--degree',
        default=3, type=int,
        help='Lagrange 有限元空间的次数, 默认为 3 次.')



parser.add_argument('--mesh_h',
        default = 0.5, type=float,
        help='网格尺寸,默认为0.5')




parser.add_argument('--show_mesh',
        default = False, type=bool,
        help='是否展示网格，默认为False.')

args = parser.parse_args()
degree = args.degree
mesh_h = args.mesh_h
show_mesh = args.show_mesh 


##########Get PDE############
PDE = Girkmann(mesh_h)
#############################




#######general mesh#########
mesh = PDE.get_mesh()
############################




s



###################Load space###################
tspace = HuZhangFiniteElementSpace(mesh, degree)
vspace = LagrangeFiniteElementSpace(mesh, degree-1, spacetype='D')
tgdof = tspace.number_of_global_dofs()
vgdof = 2*vspace.number_of_global_dofs()
gdim = 2
sh = tspace.function()
uh = vspace.function(dim=gdim)
###################get matrix####################
M = compliance_tensor_matrix(tspace,PDE.mu)
B0,B1 = tspace.parallel_div_matrix(vspace)
C = mass_vector_matrix(vspace,PDE.mu)

###################get right hand################
F0 = np.zeros(tgdof,dtype=np.float64)
F1 =  source_vector(vspace,PDE)

##################stress boundary################
isBDdof = set_essential_bc(tspace,sh,PDE)

F0 -= M@sh
F0[isBDdof] = sh[isBDdof]
F1[:,0] -= B0@sh 
F1[:,1] -= B1@sh

bdIdx = np.zeros(tgdof, dtype=int)
bdIdx[isBDdof] = 1
Tbd = spdiags(bdIdx,0,tgdof,tgdof)
T = spdiags(1-bdIdx,0,tgdof,tgdof)
M = T@M@T + Tbd
B0 = B0@T
B1 = B1@T


##################slove##########################
FF = np.r_[F0,F1.T.reshape(-1)]
AA = bmat([[M, B0.transpose(), B1.transpose()],[B0, -C, None],[B1,None,-C]],format='csr')

x = spsolve(AA,FF)


sh[:] = x[:tgdof]




Post = Post(PDE,sh,tspace)

print(Post.get_Q())
print(Post.get_M())

print(len(x))
















if show_mesh:
    fig = plt.figure()
    axes = fig.gca()
    mesh.add_plot(axes)
    mesh.find_edge(axes,showindex=True)
    #mesh.find_cell(axes,showindex=True)
    #mesh.find_node(axes,showindex=True)
    plt.show()

